package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"sync"
	"time"

	mqtt "github.com/eclipse/paho.mqtt.golang"
	"github.com/redis/go-redis/v9"
	"golang.org/x/net/context"
)

const (
	RedisHost     = "localhost"
	RedisPort     = 6379
	BrokerAddress = "localhost"
	BrokerPort    = 1883
	MaxWorkers    = 5
)

var (
	workerURLs = []string{
		"http://localhost:55043/simulate",
		"http://localhost:55049/simulate",
		"http://localhost:5000/predict",
	}
	messageQueue     = make(chan string, 100)
	redisClient      *redis.Client
	ctx              = context.Background()
	mqttClient       mqtt.Client
	mqttInitialized  = make(chan struct{})
)

type CANMessage struct {
	CANID      uint32   `json:"can_id"`
	Timestamp  int64    `json:"timestamp"`
	Data       []byte   `json:"data"`
	VehicleID  string   `json:"vehicle_id"`
	IsAnomaly  bool     `json:"is_anomaly"`
	Confidence float64  `json:"confidence,omitempty"`
}

type RequestData struct {
	UUID string                 `json:"uuid"`
	Data map[string]interface{} `json:"data"`
}

func setupRedisClient() error {
	redisAddr := fmt.Sprintf("%s:%d", RedisHost, RedisPort)
	redisClient = redis.NewClient(&redis.Options{
		Addr:     redisAddr,
		Password: "",
		DB:       0,
	})

	_, err := redisClient.Ping(ctx).Result()
	if err != nil {
		return fmt.Errorf("Failed to connect to Redis: %v", err)
	}
	log.Println("Successfully connected to Redis")
	return nil
}

func onMessageReceived(client mqtt.Client, msg mqtt.Message) {
	messageStr := string(msg.Payload())
	if len(messageStr) >= 2 && messageStr[0] == '"' && messageStr[len(messageStr)-1] == '"' {
		messageStr = messageStr[1 : len(messageStr)-1]
	}
	log.Printf("Received uuid '%s' on topic '%s', QoS %d", msg.Topic(), messageStr, msg.Qos())
	messageQueue <- messageStr
}

func setupMQTTClient() error {
	opts := mqtt.NewClientOptions()
	brokerURL := fmt.Sprintf("tcp://%s:%d", BrokerAddress, BrokerPort)
	opts.AddBroker(brokerURL)
	opts.SetClientID("LoadBalancer")
	opts.SetCleanSession(true)
	opts.SetAutoReconnect(true)
	opts.SetMaxReconnectInterval(60 * time.Second)

	var mqttInitialized = make(chan struct{})
	var mqttInitializedClosed = false
	var mqttInitializedMutex sync.Mutex

	opts.SetOnConnectHandler(func(client mqtt.Client) {
		log.Println("Connected to MQTT broker")
		if token := client.Subscribe("request", 0, onMessageReceived); token.Wait() && token.Error() != nil {
			log.Printf("Subscription failed: %v", token.Error())
		} else {
			log.Println("Subscribed to 'request' topic")
		}
		mqttInitializedMutex.Lock()
		if !mqttInitializedClosed {
			close(mqttInitialized)
			mqttInitializedClosed = true
		}
		mqttInitializedMutex.Unlock()
	})

	opts.SetConnectionLostHandler(func(client mqtt.Client, err error) {
		log.Printf("Connection lost: %v", err)
	})

	mqttClient = mqtt.NewClient(opts)
	if token := mqttClient.Connect(); token.Wait() && token.Error() != nil {
		return fmt.Errorf("Failed to connect to MQTT broker: %v", token.Error())
	}

	<-mqttInitialized

	return nil
}

func determineWorkerURL(uuid string, _ map[string]string) string {
	if len(uuid) > 8 && uuid[:8] == "vehicle:" {
		return workerURLs[2]
	}

	hash := 0
	for _, c := range uuid {
		hash += int(c)
	}
	return workerURLs[hash%2]
}

func sendRequest(url string, id int, wg *sync.WaitGroup) {
	defer wg.Done()
	log.Printf("Worker node %d started!", id)

	httpClient := &http.Client{
		Timeout: 30 * time.Second,
	}

	for {
		uuid, ok := <-messageQueue
		if !ok {
			log.Printf("Worker node %d stopped: message queue closed", id)
			return
		}

		log.Printf("Worker node %d is processing %s...", id, uuid)
		
		var dataMap map[string]interface{}
		var success bool
		
		strData, err := redisClient.Get(ctx, uuid).Result()
		if err == nil {
			if jsonErr := json.Unmarshal([]byte(strData), &dataMap); jsonErr == nil {
				success = true
				log.Printf("Worker node %d: Successfully retrieved string value from Redis", id)
			} else {
				log.Printf("Worker node %d: Error parsing Redis string: %v", id, jsonErr)
			}
		} else if err != redis.Nil {
			log.Printf("Worker node %d: Error retrieving string from Redis: %v", id, err)
		}
		
		if !success {
			redisData, err := redisClient.HGetAll(ctx, uuid).Result()
			if err == nil && len(redisData) > 0 {
				dataMap = make(map[string]interface{})
				for k, v := range redisData {
					dataMap[k] = v
				}
				success = true
				log.Printf("Worker node %d: Successfully retrieved hash value from Redis", id)
			} else if err != nil {
				log.Printf("Worker node %d: Error retrieving hash from Redis: %v", id, err)
			} else {
				log.Printf("Worker node %d: No data found in Redis for UUID %s", id, uuid)
			}
		}
		
		if !success || dataMap == nil {
			log.Printf("Worker node %d: Cannot process UUID %s, skipping", id, uuid)
			continue
		}
		
		workerURL := determineWorkerURL(uuid, nil)
		
		var jsonData []byte
		
		if workerURL == workerURLs[2] {
			jsonData, err = json.Marshal(dataMap)
		} else {
			requestData := RequestData{
				UUID: uuid,
				Data: dataMap,
			}
			jsonData, err = json.Marshal(requestData)
		}
		
		if err != nil {
			log.Printf("Worker node %d: JSON encoding error: %v", id, err)
			continue
		}
		
		req, reqErr := http.NewRequest("POST", workerURL, bytes.NewBuffer(jsonData))
		if reqErr != nil {
			log.Printf("Worker node %d: Request creation error: %v", id, reqErr)
			continue
		}
		req.Header.Set("Content-Type", "application/json")
		
		resp, respErr := httpClient.Do(req)
		if respErr != nil {
			log.Printf("Worker node %d: Request error: %v", id, respErr)
			continue
		}
		
		body, readErr := ioutil.ReadAll(resp.Body)
		resp.Body.Close()
		if readErr != nil {
			log.Printf("Worker node %d: Error reading response: %v", id, readErr)
		}
		
		if resp.StatusCode == http.StatusOK {
			log.Printf("Worker node %d received response %d, data: %s", id, resp.StatusCode, string(body))
		} else {
			log.Printf("Worker node %d received non-200 response: %d", id, resp.StatusCode)
		}
	}
}

func main() {
	if err := setupRedisClient(); err != nil {
		log.Fatalf("Failed to set up Redis: %v", err)
	}

	mqttInitialized = make(chan struct{})
	if err := setupMQTTClient(); err != nil {
		log.Fatalf("Failed to set up MQTT: %v", err)
	}

	var wg sync.WaitGroup
	for i := 0; i < MaxWorkers; i++ {
		wg.Add(1)
		go sendRequest(workerURLs[i%len(workerURLs)], i, &wg)
	}

	wg.Wait()
}
